import os
from typing import Any, Dict
from urllib.parse import parse_qsl, quote_plus

from pydantic import Field, NonNegativeInt, PositiveInt, computed_field
from pydantic_settings import BaseSettings


class DatabaseConfig(BaseSettings):
    DB_HOST: str = Field(
        description="Hostname or IP address of the database server.",
        default="localhost",
    )

    DB_PORT: PositiveInt = Field(
        description="Port number for database connection.",
        default=5432,
    )

    DB_USERNAME: str = Field(
        description="Username for database authentication.",
        default="postgres",
    )

    DB_PASSWORD: str = Field(
        description="Password for database authentication.",
        default="",
    )

    DB_DATABASE: str = Field(
        description="Name of the database to connect to.",
        default="dify",
    )

    DB_CHARSET: str = Field(
        description="Character set for database connection.",
        default="",
    )

    DB_EXTRAS: str = Field(
        description="Additional database connection parameters. Example: 'keepalives_idle=60&keepalives=1'",
        default="",
    )

    SQLALCHEMY_DATABASE_URI_SCHEME: str = Field(
        description="Database URI scheme for SQLAlchemy connection.",
        default="postgresql",
    )

    # 新增：额外的数据库配置（如规则库、工厂库等）
    RULE_DB_HOST: str = Field(default="localhost")
    RULE_DB_PORT: PositiveInt = Field(default=5432)
    RULE_DB_USERNAME: str = Field(default="postgres")
    RULE_DB_PASSWORD: str = Field(default="")
    RULE_DB_DATABASE: str = Field(default="rule_db")

    # @computed_field
    # def SQLALCHEMY_DATABASE_URI(self) -> str:
    #     db_extras = (
    #         f"{self.DB_EXTRAS}&client_encoding={self.DB_CHARSET}" if self.DB_CHARSET else self.DB_EXTRAS
    #     ).strip("&")
    #     db_extras = f"?{db_extras}" if db_extras else ""
    #     return (
    #         f"{self.SQLALCHEMY_DATABASE_URI_SCHEME}://"
    #         f"{quote_plus(self.DB_USERNAME)}:{quote_plus(self.DB_PASSWORD)}@{self.DB_HOST}:{self.DB_PORT}/{self.DB_DATABASE}"
    #         f"{db_extras}"
    #     )

    SQLALCHEMY_POOL_SIZE: NonNegativeInt = Field(
        description="Maximum number of database connections in the pool.",
        default=30,
    )

    SQLALCHEMY_MAX_OVERFLOW: NonNegativeInt = Field(
        description="Maximum number of connections that can be created beyond the pool_size.",
        default=10,
    )

    SQLALCHEMY_POOL_RECYCLE: NonNegativeInt = Field(
        description="Number of seconds after which a connection is automatically recycled.",
        default=3600,
    )

    SQLALCHEMY_POOL_PRE_PING: bool = Field(
        description="If True, enables connection pool pre-ping feature to check connections.",
        default=False,
    )

    SQLALCHEMY_ECHO: bool | str = Field(
        description="If True, SQLAlchemy will log all SQL statements.",
        default=False,
    )

    RETRIEVAL_SERVICE_EXECUTORS: NonNegativeInt = Field(
        description="Number of processes for the retrieval service, default to CPU cores.",
        default=os.cpu_count() or 1,
    )

    @computed_field
    def SQLALCHEMY_DATABASE_URI(self) -> str:
        return self._make_uri(
            scheme=self.SQLALCHEMY_DATABASE_URI_SCHEME,
            host=self.DB_HOST,
            port=self.DB_PORT,
            username=self.DB_USERNAME,
            password=self.DB_PASSWORD,
            database=self.DB_DATABASE,
            extras=self.DB_EXTRAS,
            charset=self.DB_CHARSET
        )

    def get_bind_uri(self, bind_key: str) -> str:
        mapping = {
            "rule_db": (
                self.RULE_DB_HOST,
                self.RULE_DB_PORT,
                self.RULE_DB_USERNAME,
                self.RULE_DB_PASSWORD,
                self.RULE_DB_DATABASE,
            ),
        }
        if bind_key not in mapping:
            raise ValueError(f"Unknown bind key: {bind_key}")

        host, port, user, pwd, db = mapping[bind_key]
        return self._make_uri(
            scheme=self.SQLALCHEMY_DATABASE_URI_SCHEME,
            host=host,
            port=port,
            username=user,
            password=pwd,
            database=db,
            extras=self.DB_EXTRAS,
            charset=self.DB_CHARSET
        )

    @staticmethod
    def _make_uri(scheme: str, host: str, port: int, username: str, password: str, database: str, extras: str,
                  charset: str) -> str:
        db_extras = f"{extras}&client_encoding={charset}" if charset else extras
        db_extras = f"?{db_extras.strip('&')}" if db_extras.strip("&") else ""
        return (
            f"{scheme}://{quote_plus(username)}:{quote_plus(password)}"
            f"@{host}:{port}/{database}{db_extras}"
        )

    @computed_field  # type: ignore[misc]
    def SQLALCHEMY_BINDS(self) -> Dict[str, str]:
        return {
            "rule_db": self.get_bind_uri("rule_db"),
            # 可继续扩展其他库
        }

    @computed_field  # type: ignore[misc]
    @property
    def SQLALCHEMY_ENGINE_OPTIONS(self) -> dict[str, Any]:
        # Parse DB_EXTRAS for 'options'
        db_extras_dict = dict(parse_qsl(self.DB_EXTRAS))
        options = db_extras_dict.get("options", "")
        # Always include timezone
        timezone_opt = "-c timezone=UTC"
        if options:
            # Merge user options and timezone
            merged_options = f"{options} {timezone_opt}"
        else:
            merged_options = timezone_opt

        connect_args = {"options": merged_options}

        return {
            "pool_size": self.SQLALCHEMY_POOL_SIZE,
            "max_overflow": self.SQLALCHEMY_MAX_OVERFLOW,
            "pool_recycle": self.SQLALCHEMY_POOL_RECYCLE,
            "pool_pre_ping": self.SQLALCHEMY_POOL_PRE_PING,
            "connect_args": connect_args,
        }
